"""
Koopman算子网络
作者: magical857
日期: 2025-11-01
"""

import torch
import torch.nn as nn
from .lstm_autoencoder import LSTMAutoencoder


class KoopmanOperator(nn.Module):
    """Koopman算子 - 在高维空间中的线性演化算子"""
    
    def __init__(self, koopman_dim):
        super(KoopmanOperator, self).__init__()
        
        # Koopman算子矩阵 K
        self.K = nn.Linear(koopman_dim, koopman_dim, bias=False)
        
        # 可选: 添加正则化以保持算子的稳定性
        # 这里使用简单的线性层,也可以添加约束使其接近正交矩阵
    
    def forward(self, psi):
        """
        应用Koopman算子: psi_{t+1} = K * psi_t
        
        Args:
            psi: (batch_size, koopman_dim) - 当前时刻Koopman嵌入
        Returns:
            psi_next: (batch_size, koopman_dim) - 下一时刻Koopman嵌入
        """
        return self.K(psi)
    
    def multi_step_prediction(self, psi, num_steps):
        """
        多步预测
        
        Args:
            psi: (batch_size, koopman_dim)
            num_steps: 预测步数
        Returns:
            predictions: (batch_size, num_steps, koopman_dim)
        """
        predictions = []
        psi_current = psi
        
        for _ in range(num_steps):
            psi_next = self.forward(psi_current)
            predictions.append(psi_next)
            psi_current = psi_next
        
        return torch.stack(predictions, dim=1)


class KoopmanNetwork(nn.Module):
    """
    完整的Koopman网络: LSTM-AE + Koopman算子
    """
    
    def __init__(self, state_dim, control_dim, hidden_dim, num_layers, 
                 koopman_dim, dropout=0.2):
        super(KoopmanNetwork, self).__init__()
        
        self.state_dim = state_dim
        self.control_dim = control_dim
        self.koopman_dim = koopman_dim
        
        # LSTM自编码器
        self.autoencoder = LSTMAutoencoder(
            state_dim=state_dim,
            control_dim=control_dim,
            hidden_dim=hidden_dim,
            num_layers=num_layers,
            koopman_dim=koopman_dim,
            dropout=dropout
        )
        
        # Koopman算子
        self.koopman_operator = KoopmanOperator(koopman_dim)
    
    def forward(self, states, controls):
        """
        前向传播: 重构 + Koopman预测
        
        Args:
            states: (batch_size, seq_len, state_dim)
            controls: (batch_size, seq_len, control_dim)
        Returns:
            reconstruction: 重构的状态
            koopman_embedding: Koopman嵌入
            koopman_prediction: Koopman空间中的预测
        """
        # 自编码器重构
        reconstruction, koopman_embedding = self.autoencoder(states, controls)
        
        # Koopman算子预测 (在序列的每个时间步应用)
        batch_size, seq_len, koopman_dim = koopman_embedding.shape
        
        if seq_len > 1:
            # 取前 seq_len-1 个时间步
            psi_t = koopman_embedding[:, :-1, :]  # (batch, seq_len-1, koopman_dim)
            psi_t_flat = psi_t.reshape(-1, koopman_dim)
            
            # 应用Koopman算子
            psi_next_flat = self.koopman_operator(psi_t_flat)
            psi_next = psi_next_flat.reshape(batch_size, seq_len-1, koopman_dim)
            
            koopman_prediction = psi_next
        else:
            koopman_prediction = None
        
        return reconstruction, koopman_embedding, koopman_prediction
    
    def predict_future(self, states, controls, num_steps):
        """
        长期预测
        
        Args:
            states: (batch_size, seq_len, state_dim) - 初始序列
            controls: (batch_size, num_steps, control_dim) - 未来控制
            num_steps: 预测步数
        Returns:
            predictions: (batch_size, num_steps, state_dim)
        """
        # 编码初始状态到Koopman空间
        koopman_emb = self.autoencoder.encode(states, controls[:, :states.shape[1], :])
        
        # 取最后一个时间步的嵌入
        psi_0 = koopman_emb[:, -1, :]  # (batch_size, koopman_dim)
        
        # 多步Koopman预测
        psi_predictions = self.koopman_operator.multi_step_prediction(psi_0, num_steps)
        # psi_predictions: (batch_size, num_steps, koopman_dim)
        
        # 解码回状态空间
        predictions = self.autoencoder.decode(psi_predictions)
        
        return predictions
    
    def get_koopman_matrix(self):
        """获取Koopman算子矩阵"""
        return self.koopman_operator.K.weight.data.cpu().numpy()


if __name__ == "__main__":
    # 测试
    batch_size = 8
    seq_len = 50
    state_dim = 4
    control_dim = 2
    
    model = KoopmanNetwork(
        state_dim=state_dim,
        control_dim=control_dim,
        hidden_dim=128,
        num_layers=2,
        koopman_dim=64,
        dropout=0.2
    )
    
    states = torch.randn(batch_size, seq_len, state_dim)
    controls = torch.randn(batch_size, seq_len, control_dim)
    
    reconstruction, koopman_emb, koopman_pred = model(states, controls)
    
    print(f"重构形状: {reconstruction.shape}")
    print(f"Koopman嵌入形状: {koopman_emb.shape}")
    if koopman_pred is not None:
        print(f"Koopman预测形状: {koopman_pred.shape}")
    
    # 测试长期预测
    future_controls = torch.randn(batch_size, 20, control_dim)
    future_predictions = model.predict_future(states, 
                                             torch.cat([controls, future_controls], dim=1), 
                                             num_steps=20)
    print(f"长期预测形状: {future_predictions.shape}")
    
    print(f"总参数量: {sum(p.numel() for p in model.parameters()):,}")